# @Author：cnzdy
# @Email：cn_zdy@126.com
# @Time: 2021/10/9 11:05
# @File: trainer.py
import numpy as np
import torch
from torchnet import meter
from tqdm import tqdm

from pytorch.classifier.options import Options
from tutorial.pytorch.classifier.utils import AvgrageMeter
from tutorial.pytorch.classifier.analysis import plot_training, log
# from utils_viz.viz_visdom import viz_visdom

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


class Trainer(object):
    def __init__(self):
        self.args = Options
        self.args.model.to(device)

        self.optimizer = torch.optim.SGD(self.args.model.parameters(), lr=self.args.lr,
                                         momentum=0.9)
        self.loss = torch.nn.CrossEntropyLoss()

        self.loss_meter = meter.AverageValueMeter()
        self.confusion_meter = meter.ConfusionMeter(Options.num_classes)
        self.top1 = AvgrageMeter()

        self.previous_loss = 1e10
        # 随便设置一个比较大的数，用于保存精度最高的模型
        self.min_loss = 100000

        # 用于绘制loss和acc曲线
        self.train_loss = []
        self.valid_loss = []
        self.train_acc = []
        self.valid_acc = []

    @plot_training(Options.analysis_path)
    def run(self):
        for e in range(self.args.epochs):
            self.init_epoch(e)
            self.train()
            self.validate()
            self.save()
            self.adjust_lr()

        return self.train_loss, self.valid_loss, self.train_acc, self.valid_acc

    def init_epoch(self, e):
        # epoch计数
        self.e = e
        # 把模型恢复为训练模式
        self.args.model.train()
        self.loss_meter.reset()
        self.confusion_meter.reset()
        self.top1.reset()

    # @log
    # @viz_visdom
    def train(self):
        i = 0
        loss_sum = 0.0
        loss_avg = 0.0
        acc = 0.0

        print('train model...')
        # 转换成tqdm类型 以方便增加日志的输出
        self.args.train_loader = tqdm(self.args.train_loader)
        self.args.train_loader.set_description(f'Epoch:{self.e + 1:04d}/{self.args.epochs:04d} '
                                               f'lr: {self.args.lr}')

        # enumerate的第二个参数0表示：从索引从0开始，假如为1的话，那索引就从1开始。
        for i, data in enumerate(self.args.train_loader, 0):
            # data里面包含图像数据（inputs）(tensor类型的）和标签（labels）(tensor类型）。
            # 用方法to()可以将Tensor在CPU和GPU（需要硬件支持）之间相互移动;
            # to()还可以同时更改数据类型
            input_batch = data[0].to(device, non_blocking=True)
            label_batch = data[1].to(device, non_blocking=True)
            # print(f"input batch shape: {input_batch.shape}")
            # print(f"label batch: {label_batch}")

            # 梯度清零
            self.optimizer.zero_grad()
            output_batch = self.args.model(input_batch)
            loss = self.loss(output_batch, label_batch)
            loss.backward()

            # 所有optimizer都实现了step()方法，调用这个方法可以更新参数
            # 每次用backward()这类方法计算出了梯度后，就可以调用一次这个方法来更新参数。
            self.optimizer.step()

            # meters update and visualize  更新统计指标及可视化
            self.loss_meter.add(loss.item())

            # detach 一下更安全保险
            self.confusion_meter.add(output_batch.detach(), label_batch.detach())
            loss_sum += loss.item()
            loss_avg = loss_sum / (i + 1)

            # 计算模型在验证集上的准确率等信息(与validate()函数功能一样）
            # 先将所有验证集样本的预测结果保存到ConfusionMeter中，然后通过混淆矩阵计算整个验证集的acc
            cm_value = self.confusion_meter.value()
            acc = 100. * (cm_value[0][0] + cm_value[1][1]) / (cm_value.sum())

            postfix = {'train_loss': '%.6f' % loss_avg, 'train_acc': '%.6f' % acc}
            self.args.train_loader.set_postfix(log=postfix)

        self.train_loss.append(loss_avg)
        self.train_acc.append(acc)

        # 用于@log保存日志，@viz_visdom显示
        result = {"epoch": self.e,
                  "lr": self.args.lr,
                  "loss": loss_avg,
                  "acc": acc,
                  "cm": cm_value,
                  "optimizer": self.optimizer}
        return result

    def adjust_lr(self):
        # update learning rate  如果损失不再下降，则降低学习率
        if self.loss_meter.value()[0] > self.previous_loss:
            self.args.lr = self.args.lr * self.args.lr_decay
            # 第二种降低学习率的方法:不会有moment等信息的丢失
            for param_group in self.optimizer.param_groups:
                param_group['lr'] = self.args.lr

        self.previous_loss = self.loss_meter.value()[0]

    def save(self):
        # 每次只保存最好的模型
        if self.valid_loss[self.e] < self.min_loss:
            self.min_loss = self.valid_loss[self.e]
            print("save model...\n")

            name = self.args.models_path + self.args.project_name + '_' + \
                   self.args.model_name + '_best.pth'
            torch.save(self.args.model.state_dict(), name)
            return name

    def validate(self):
        """训练完一轮以后，用验证集验证模型准确率.
           验证时显示每个batch的loss和acc，最后给出整个验证集的平均acc.
        """
        print("validate model...")

        with torch.no_grad():  # 进行评测的时候网络不更新梯度
            self.args.model.eval()
            val_top1 = AvgrageMeter()
            validate_loader = tqdm(self.args.val_loader)
            loss_sum = 0.0

            for i, data in enumerate(validate_loader, 0):  # 0是下标起始位置默认为0
                inputs, labels = data[0].to(device, non_blocking=True), \
                                 data[1].to(device, non_blocking=True)
                outputs = self.args.model(inputs)
                loss = self.loss(outputs, labels)
                loss_sum += loss.item()

                # 计算验证集 acc
                prec1, prec2 = self.accuracy(outputs, labels, topk=(1, 2))
                n = inputs.size(0)
                val_top1.update(prec1.item(), n)

                postfix = {'valid_loss': '%.6f' % (loss_sum / (i + 1)),
                           'valid_acc': '%.6f' % val_top1.avg}
                validate_loader.set_postfix(log=postfix)

            print(f"val_acc: {val_top1.avg}")
            self.valid_loss.append(loss_sum / (i + 1))
            self.valid_acc.append(val_top1.avg)

    def accuracy(self, output, label, topk=(1,)):
        """topk的准确率计算"""
        maxk = max(topk)
        batch_size = label.size(0)

        # torch.topk(input, k, dim=None, largest=True, sorted=True, out=None) -> (Tensor, LongTensor)
        # 返回某一维度前k个的索引
        # input：一个tensor数据
        # k：指明是得到前k个数据以及其index
        # dim： 指定在哪个维度上排序， 默认是最后一个维度
        # largest：如果为True，按照大到小排序； 如果为False，按照小到大排序
        # sorted：返回的结果按照顺序返回
        # out：可缺省，不要
        # 获取前K的索引
        _, pred = output.topk(maxk, 1, True, True)  # 使用topk来获得前k个的索引
        pred = pred.t()  # 进行转置
        # eq按照对应元素进行比较 view(1,-1) 自动转换到行为1,的形状， expand_as(pred) 扩展到pred的shape
        # expand_as 执行按行复制来扩展，要保证列相等
        # 与正确标签序列形成的矩阵相比，生成True/False矩阵
        correct = pred.eq(label.view(1, -1).expand_as(pred))
        #     print(correct)

        rtn = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0)  # 前k行的数据 然后平整到1维度，来计算true的总个数
            rtn.append(correct_k.mul_(100.0 / batch_size))  # mul_() ternsor 的乘法  正确的数目/总的数目 乘以100 变成百分比
        return rtn
